#%%
### This script was written by Eric Deal, 06.22.2020
### Script should be run with python 3
import numpy as np, pandas as pd, pathlib, seaborn as sns, matplotlib.pyplot as plt, scipy.special as ss

# Grain dimension data
# The grain types we used consisted of two kinds - grains with identical sizes
# and shapes, and those with variation from grain to grain.

# create empty dictionaries to be filled with grain dimension measurements
grains, hole_props = {}, {}
A, dA, B, dB, C, dC = {}, {}, {}, {}, {}, {}
A_vec, B_vec, C_vec = {}, {}, {}
dn, ddn = {}, {}

### First: the grains with identical shapes
# In this case, we were able to characterize the shape of the grain to within
# measurement error - about 0.1 mm. We measured the longest axis - denoted 'A'
# in this script, and then the largest axis perpendicular to 'A' - referred to
# as 'B' here. Finally, axis 'C' is that perpendicular to both 'A' and 'B'.

# Glass beads ('gb')
# The glass beads are spheres, and can be characterized by a single diameter.
# We still use the A, B, C terminology here to ease comparison
grains['gb'] = 'gb' # name grains
A['gb'], B['gb'], C['gb'] = 0.0049 * np.ones(3) # grain diameter
dA['gb'], dB['gb'], dC['gb'] = 5e-5 * np.ones(3) # measurement error on grain diameter

# Faceted ellipsoids ('oc')
# The faceted ellipsoids are circular in a cross-section perpendicular to their
# shortest axis, but we stick with the three axes as with the glass beads. They
# also have a hole in them, which we include measurements on to help with
# measuring their density.
grains['oc'] = 'oc' # name grains
hole_props['oc'] = {'hole_l': 5e-3, 'dhole_l': 1e-4, 'hole_r': 5e-4, 'dhole_r': 1e-4} # hole length, hole length error, hole radius, hole radius error
A['oc'], B['oc'] = 0.006 * np.ones(2) # grain diameter
C['oc'] = 0.005 # grain diameter in shortest direction
dA['oc'], dB['oc'], dC['oc'] = 1e-4 * np.ones(3) # measurement error on grain diameter

### Second: the grains with variable shapes
# In this case, we were able to characterize the shape of the grain as an average
# of grain dimensions across a number of grains. In the case of the natural river
# gravel, we make use of ca. 1600 that were put into a microCT scanner, giving
# very high resolution shape data. In the case of the final two granular materials
# the gravel-like glass chips and the glass cubic prisms, we measured the three
# perpendicular axes on a number of grains by hand.

# Natural river gravel ('ng')
# River gravel sourced from the Fraser river delta and sieved to a narrow grain size distribution
grains['ng'] = 'ng' # name grains
# load grain shape data from CT scan results
os.getcwd()
df_ng_shape = pd.DataFrame(np.hstack([np.load(path) for path in list(pathlib.Path('observations/2_data_shape').glob('*Raw*.npy'))]).T)
df_ng_shape.columns = ['Grain Labels', 'Grain Elongation', 'MOS5 Elongation', 'MOS10 Elongation', 'MOS20 Elongation', 'Grain Flatness', 'MOS5 Flatness', 'MOS10 Flatness', 'MOS20 Flatness', 'Grain Corey Shape Factor', 'MOS5 Corey Shape Factor', 'MOS10 Corey Shape Factor', 'MOS20 Corey Shape Factor', 'Grain Krumbein Intercept Sphericity', 'MOS5 Krumbein Intercept Sphericity', 'MOS10 Krumbein Intercept Sphericity', 'MOS20 Krumbein Intercept Sphericity', 'Grain Sneed and Folk Max Projected Sphericity', 'MOS5 Sneed and Folk Max Projected Sphericity', 'MOS10 Sneed and Folk Max Projected Sphericity', 'MOS20 Sneed and Folk Max Projected Sphericity', 'Grain Wilson and Huang Shape Factor', 'MOS5 Wilson and Huang Shape Factor', 'MOS10 Wilson and Huang Shape Factor', 'MOS20 Wilson and Huang Shape Factor', 'Grain Riley Circularity', 'MOS5 Riley Circularity', 'MOS10 Riley Circularity', 'MOS20 Riley Circularity', 'Grain Formfactor L (mm)', 'MOS5 Formfactor L', 'MOS10 Formfactor L', 'MOS20 Formfactor L', 'Grain Formfactor i (mm)', 'MOS5 Formfactor i', 'MOS10 Formfactor i', 'MOS20 Formfactor i', 'Grain Formfactor S (mm)', 'MOS5 Formfactor S', 'MOS10 Formfactor S', 'MOS20 Formfactor S', 'Grain Volume (mm3)', 'MOS5 Volume', 'MOS10 Volume', 'MOS20 Volume']
A_vec['ng'], B_vec['ng'], C_vec['ng'] = [df_ng_shape['Grain Formfactor %s (mm)' % char]*1e-3 for char in ['L', 'i', 'S']] # vector of all measured A, B and C axes
A['ng'], B['ng'], C['ng'] = [df_ng_shape['Grain Formfactor %s (mm)' % char].mean()*1e-3 for char in ['L', 'i', 'S']] # average of all measured A, B and C axes
dA['ng'], dB['ng'], dC['ng'] = [df_ng_shape['Grain Formfactor %s (mm)' % char].std()*1e-3/np.sqrt(df_ng_shape['Grain Formfactor %s (mm)' % char].size) for char in ['L', 'i', 'S']] # standard deviation of all measured A, B and C axes as a measure of error/variability

# Cubic prisms ('ls')
# Grains with a square cross-section perpendicular to the longest axis. The
# variation in grain size lies mostly in the length of the longest axis. The
# grains have a hole through the center which is also measured
grains['ls'] = 'ls' # name grains
# hand-made grain size measurements of three perpendicular axes.
ABC = np.array([[3.5, 3.5, 4],[3.5, 4, 6],[3, 3, 4],[3.5, 3.5, 4],[3.5, 3.5, 5],[3.5, 3.5, 5],[3, 3, 3.5],[3.3, 3.3, 3.5],[3.5, 3.5, 4],[3.5, 3.5, 5]])
hole_props['ls'] = {'hole_l': ABC[:,2].mean() / 1e3, 'dhole_l': ABC[:,2].std()/np.sqrt(ABC.shape[0]) / 1e3, 'hole_r': 5e-4, 'dhole_r': 1e-4} # hole length, hole length error, hole radius, hole radius error
A_vec['ls'], B_vec['ls'], C_vec['ls'] = [np.array([func(shape) for shape in ABC])*1e-3 for func in [lambda x: np.max(x), lambda x: np.sort(x)[1], lambda x: np.min(x)]] # vector of all measured A, B and C axes
A['ls'], B['ls'], C['ls'] = [np.array([func(shape) for shape in ABC]).mean()*1e-3 for func in [lambda x: np.max(x), lambda x: np.sort(x)[1], lambda x: np.min(x)]] # average of all measured A, B and C axes
dA['ls'], dB['ls'], dC['ls'] = [np.array([func(shape) for shape in ABC]).std()*1e-3/np.sqrt(np.array([func(shape) for shape in ABC]).size) for func in [lambda x: np.max(x), lambda x: np.sort(x)[1], lambda x: np.min(x)]] # standard deviation of all measured A, B and C axes as a measure of error/variability

# Glass chips ('gg')
# Grains consisting of glass chips with random shapes. These grains were measured
# by hand. They also have a hole which was measured as well.
grains['gg'] = 'gg' # name grains
# hand-made grain size measurements of three perpendicular axes.
ABC = np.vstack([
        np.array([[5.5, 7, 13.5],[3, 5, 6],[3, 7, 7],[3, 7, 9],[5.5, 6.5, 9],[5.5, 7, 11.5],[3.5, 4.5, 8],[3, 5.5, 7],[4, 6.5, 9],[3.5, 6.5, 8.5],[3.5, 6, 7.5],[3, 5, 7]]),
        np.array([[3, 5, 8], [5, 7, 13], [5, 6, 11], [4, 5, 12], [3, 4, 7], [2, 6, 12], [3, 7, 8], [4, 7, 10], [2, 5, 7], [5, 6, 7], [3.5, 6, 10], [2.5, 5, 6], [2, 6, 10], [4, 4, 10], [3, 5, 10], [3, 4, 5], [3, 4, 6], [3.5, 5, 7]])
        ])
hole_props['gg'] = {'hole_l': ABC[:,1].mean() / 1e3, 'dhole_l': ABC[:,1].std()/np.sqrt(ABC.shape[0]) / 1e3, 'hole_r': 5e-4, 'dhole_r': 1e-4} # hole length, hole length error, hole radius, hole radius error
A_vec['gg'], B_vec['gg'], C_vec['gg'] = [np.array([func(shape) for shape in ABC])*1e-3 for func in [lambda x: np.max(x), lambda x: np.sort(x)[1], lambda x: np.min(x)]] # vector of all measured A, B and C axes
A['gg'], B['gg'], C['gg'] = [np.array([func(shape) for shape in ABC]).mean()*1e-3 for func in [lambda x: np.max(x), lambda x: np.sort(x)[1], lambda x: np.min(x)]] # average of all measured A, B and C axes
dA['gg'], dB['gg'], dC['gg'] = [np.array([func(shape) for shape in ABC]).std()*1e-3/np.sqrt(np.array([func(shape) for shape in ABC]).size) for func in [lambda x: np.max(x), lambda x: np.sort(x)[1], lambda x: np.min(x)]] # standard deviation of all measured A, B and C axes as a measure of error/variability

### Diameter of volume-equivalent sphere
# Load particle volume data from density measurements
df_vol = pd.read_csv('1_var_density.csv').set_index('Unnamed: 0').rename_axis(None)
df_vol.grain_vol['ng'], df_vol.dgrain_vol['ng'] = df_ng_shape['Grain Volume (mm3)'].mean()*1e-9, df_ng_shape['Grain Volume (mm3)'].std()*1e-9
dn = {grain: (6*df_vol.grain_vol[grain]/np.pi)**(1/3) for grain in grains} # use the average grain volume to find the diameter of the volume equivalent sphere
ddn = {grain: 2*df_vol.dgrain_vol[grain]/(np.pi*(6*df_vol.grain_vol[grain]/np.pi)**(2/3)) for grain in grains} # Use gaussian error propogation to estimate error in dn
ddn['gg'] *= 10 # multiply the error by 10 for the glass gravel to better match variation seen in A, B and C for that material.

### Corey Shape Factor (C/sqrt(AB))
CSF_vec = {grain: C_vec[grain]/np.sqrt(A_vec[grain]*B_vec[grain]) for grain in ['ng', 'ls', 'gg']} # Corey shape factor for grain materials with variable shapes
CSF = {grain: CSF_vec[grain].mean() for grain in ['ng', 'ls', 'gg']} # Mean Corey shape factor for these grain types
dCSF = {grain: CSF_vec[grain].std()/np.sqrt(CSF_vec[grain].size) for grain in ['ng', 'ls', 'gg']} # standard deviation as a measure of the variation in grain shape
CSF.update({grain: C[grain]/np.sqrt(A[grain]*B[grain]) for grain in ['oc', 'gb']}) # Corey shape factor for grain types with consistent shapes
dCSF.update({grain: np.sqrt((dC[grain]/np.sqrt(A[grain]*B[grain]))**2 + (B[grain]*C[grain]*dA[grain]/(2*(A[grain]*B[grain])**(3/2)))**2 + (A[grain]*C[grain]*dB[grain]/(2*(A[grain]*B[grain])**(3/2)))**2) for grain in ['oc', 'gb']}) # Gaussian error propogration to estimate error in Corey shape factor for grain types with consistent shapes
n = {grain: CSF_vec[grain].size for grain in ['ng', 'ls', 'gg']}
n.update({grain: 1 for grain in ['oc', 'gb']}) 

# CSF_vec = {grain: C_vec[grain]/np.sqrt(A_vec[grain]*B_vec[grain]) for grain in ['ng', 'ls', 'gg']} # Corey shape factor for grain materials with variable shapes


# combine all the data into a single pandas dataframe
df = pd.DataFrame({'A': A, 'dA': dA, 'B': B, 'dB': dB, 'C': C, 'dC': dC, 'dn': dn, 'ddn': ddn, 'CSF': CSF, 'dCSF': dCSF, 'A_vec': A_vec, 'B_vec': B_vec, 'C_vec': C_vec, 'hole_props': hole_props, 'n': n})
# save data to .csv file.
df.to_csv('2_var_shape.csv')

#%%
def ellipse_perm(a,b): return 4*a*ss.ellipe(1.0 - b**2/a**2)  
def ellipse_area(a,b): return np.pi * a * b
def ellipse_circ(a,b): return 4*np.pi*ellipse_area(a,b)/ellipse_perm(a,b)**2

Circ_vec = {grain: ellipse_circ(C_vec[grain], A_vec[grain]) for grain in ['ng', 'ls', 'gg']} # Corey shape factor for grain materials with variable shapes
Circ = {grain: Circ_vec[grain].mean() for grain in ['ng', 'ls', 'gg']} # Mean Corey shape factor for these grain types
dCirc = {grain: Circ_vec[grain].std()/np.sqrt(Circ_vec[grain].size) for grain in ['ng', 'ls', 'gg']}
Circ.update({grain: ellipse_circ(C[grain], B[grain]) for grain in ['oc', 'gb']}) # Corey shape factor for grain types with consistent shapes

#%%
fig, axss = plt.subplots(2,3, figsize=(9, 6), sharey=False)
axs = axss[0]
axs2 = axss[1]

nbins = {'gg': 12, 'ng': 31, 'ls': 6}
rbins = {'gg': (2,14), 'ng': (2,10), 'ls': (2,8)}
grain_name = {'gg': 'Rounded chips', 'ng': 'Natural gravel', 'ls': 'Rectangular prisms'}
for i, grain in enumerate(['gg', 'ng', 'ls']):
        axs[i].hist(np.array(A_vec[grain]).astype(float)*1e3, density=True, label='a', alpha=.3, bins=nbins[grain], range=rbins[grain], histtype='step', color='k', fill=True) # A
        axs[i].hist(np.array(B_vec[grain]).astype(float)*1e3, density=True, label='b', alpha=.3, bins=nbins[grain], range=rbins[grain], histtype='step', fill=True) # B
        axs[i].hist(np.array(C_vec[grain]).astype(float)*1e3, density=True, label='c', alpha=.3, bins=nbins[grain], range=rbins[grain], histtype='step', fill=True) # C
        axs[i].set_title(grain_name[grain])

ylims = [15, 15, 15]
for i in range(3):
        axs[i].set_xlim(0,ylims[i])
        axs[i].set_ylim(0,1.5)
        axs[i].legend()
        axs[i].set_xlabel('Diameter (mm)')
axs[0].set_ylabel('Relative freq.')

nbins = {'gg': 8, 'ng': 30, 'ls': 2}
rbins = {'gg': (0,1), 'ng': (0,1), 'ls': (0,1)}
for i, grain in enumerate(['gg', 'ng', 'ls']):
        CSF = np.array(C_vec[grain]).astype(float)/np.sqrt(np.array(A_vec[grain]).astype(float)*np.array(B_vec[grain]).astype(float))
        print(CSF)
        axs2[i].hist(np.array(CSF), density=True, alpha=.5, bins=nbins[grain], histtype='step', fill=True, color='k', label=r'$S_f$')
        axs2[i].set_title('n = %g' % A_vec[grain].size)

ylims = [1, 1, 1]
for i in range(3):
        axs2[i].set_xlim(0,ylims[i])
        axs2[i].set_ylim(0,5)
        # axs2[i].legend()
        axs2[i].set_xlabel('Corey Shape Factor')
axs2[0].set_ylabel('Relative freq.')
plt.tight_layout()

axs[0].text(0.85, 1.35, 'a)')
axs[1].text(0.85, 1.35, 'b)')
axs[2].text(0.85, 1.35, 'c)')
axs2[0].text(0.05, 4.5, 'd)')
axs2[1].text(0.05, 4.5, 'e)')
axs2[2].text(0.05, 4.5, 'f)')
plt.savefig('Grain_size_dists.png')

